import numpy as np
import quimb as qu
import quimb.tensor as qtn
from quimb.tensor import Tensor, TensorNetwork

I, X, Y, Z = np.eye(2), np.array([[0.0, 1], [1, 0]]), np.array([[0.0, -1j], [1j, 0]]), np.array([[1.0, 0], [0, -1]])

I_o1 = np.array([np.sqrt(2), 0, 0, 0])
I_o2 = np.identity(4)
Xp_o1 = np.array([1, 1, 0, 0]) / np.sqrt(2)
Xm_o1 = np.array([1, -1, 0, 0]) / np.sqrt(2)
Yp_o1 = np.array([1, 0, 1, 0]) / np.sqrt(2)
Ym_o1 = np.array([1, 0, -1, 0]) / np.sqrt(2)
Zp_o1 = np.array([1, 0, 0, 1]) / np.sqrt(2)
Zm_o1 = np.array([1, 0, 0, -1]) / np.sqrt(2)

meas_map = {
	'x+': Xp_o1, 'x-': Xm_o1,
	'y+': Yp_o1, 'y-': Ym_o1,
	'z+': Zp_o1, 'z-': Zm_o1,
}
"""Setting"""
N_qubit, N_set, N_shot, step = 10, 1, 4000, 7
J, h, dt = 0.5236, 1.0, 0.5
probs = [0, 0, 1]
max_bond, cutoff = 800, 1e-10

"""MPO representation of trotter step"""
import pathlib

reduced_noisy_trotter_step_filename = f"cache/reduced_noisy_trotter_step, step={step}, J={J:.4f}, h={h:.2f}, dt={dt:.2f}, max_bond={max_bond}, cutoff={cutoff:.1e}"
if not pathlib.Path(reduced_noisy_trotter_step_filename).exists(): raise FileNotFoundError()
rhoT_mps = qu.load_from_disk(reduced_noisy_trotter_step_filename)

"""Measurement"""
import tqdm

meas_settings = np.random.choice(['x', 'y', 'z'], size=(N_set, N_qubit), p=probs)
datas = []
for i, meas_setting in enumerate(meas_settings):
	outcomes = []
	pbar = tqdm.trange(N_shot, desc=f"Measuring {i + 1}/{N_set}", leave=True)
	for j in pbar:
		meas = TensorNetwork([Tensor(I_o1, inds=(f"k{_}",), tags=[f"I{_}"]) for _ in range(N_qubit)], virtual=True)
		outcome, prob = [], 1.0
		for k, mea_setting in enumerate(meas_setting):
			meas.pop_tensor(k)
			meas_p = Tensor(meas_map[mea_setting + '+'], inds=(f"k{k}",), tags=[f"I{k}"])
			meas_m = Tensor(meas_map[mea_setting + '-'], inds=(f"k{k}",), tags=[f"I{k}"])
			prob_p = np.real(rhoT_mps @ (meas | meas_p))
			prob_m = np.real(rhoT_mps @ (meas | meas_m))
			# prob = prob_p + prob_m
			pbar.set_postfix({'prob': (prob_p + prob_m) / prob})
			if np.random.uniform(0, prob_p + prob_m) < prob_p:
				meas |= meas_p
				prob = prob_p
				outcome.append(mea_setting + '+')
			else:
				meas |= meas_m
				prob = prob_m
				outcome.append(mea_setting + '-')
		outcomes.append(outcome)
	datas.append(outcomes)

"""Saving the data"""
import json

data_filename = f"data/Outcomes_reduced, N_qubit={N_qubit}, N_set={N_set}, N_shot={N_shot}, step={step}, max_bond={max_bond}.json"
file = pathlib.Path(data_filename)

data_dict = {
	"N_qubit" : N_qubit,
	"N_set"   : N_set,
	"N_shot"  : N_shot,
	"steps"   : step,
	"max_bond": max_bond,
	"cutoff"  : cutoff,
	"probs"   : probs,
	"datas"   : [
		{
			'meas'    : meas_setting.tolist(),
			'outcomes': outcomes
		}
		for meas_setting, outcomes in zip(meas_settings, datas)
	]
}

with open(file, 'w') as f:
	json.dump(data_dict, f)  # , indent=4)
	print(f"Saving Measurement data to File {data_filename}")
